# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# Licensed under the 【火山方舟】原型应用软件自用许可协议
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     https://www.volcengine.com/docs/82379/1433703
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import asyncio
from abc import ABC
from typing import List, Optional
from http import HTTPStatus
from dashscope import Application

# from volcenginesdkarkruntime import AsyncArk
# from volcenginesdkarkruntime.types.bot_chat import BotChatCompletion

#from .search_engine import SearchEngine, SearchResult, SearchReference
from .search_engine import SearchEngine, SearchResult, SearchReference

"""
using bailian(百炼) bot (with search plugin) to search
"""


class BailianBotSearchEngine(SearchEngine, ABC):
    api_key: Optional[str] = None

    def __init__(
            self,
            bot_id: str,
            api_key: Optional[str] = None,
    ):
        super().__init__()
        self._bot_id = bot_id
        self.api_key = api_key

    def search(self, queries: List[str]) -> List[SearchResult]:
        return asyncio.run(self.asearch(queries=queries))

    async def asearch(self, queries: List[str]) -> List[SearchResult]:
        tasks = [self._single_search(query) for query in queries]
        task_results = await asyncio.gather(*tasks)
        return [
            result for result in task_results
        ]

    async def _single_search(self, query: str) -> SearchResult:
        response = await self._run_bot_search(query)
        return self._format_result(response, query)

    async def _run_bot_search(self, query: str):
        max_retries = 3
        retry_count = 0
        while retry_count < max_retries:
            try:
                response = Application.call(
                    # 若没有配置环境变量，可用百炼API Key将下行替换为：api_key="sk-xxx"。但不建议在生产环境中直接将API Key硬编码到代码中，以减少API Key泄露风险。
                    api_key="sk-cacafff2c8d344be6702adf4f",
                    app_id='dbce22d58d7b4c7b74d357b362b',# 替换为实际的应用 ID
                    prompt=query,
                    timeout=160)

                if response.status_code != HTTPStatus.OK:
                    print(f'request_id={response.request_id}')
                    print(f'code={response.status_code}')
                    print(f'message={response.message}')
                    print(f'请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code')
                    if retry_count < max_retries - 1:
                        retry_count += 1
                        await asyncio.sleep(1)
                        continue
                    raise ValueError(f"API request failed: {response.message}")
                break
            except Exception as e:
                if retry_count < max_retries - 1:
                    retry_count += 1
                    await asyncio.sleep(1)
                    continue
                raise
        
        if not response.output or not hasattr(response.output, 'text'):
            raise ValueError("Invalid response format from API")
            
        #print(response.output.text)
        return response.output

    @classmethod
    def _format_result(cls, response, query: str) -> SearchResult:
        if not response or not hasattr(response, 'text') or not response.text:
            raise ValueError("Invalid response content")
            
        return SearchResult(
            query=query,
            summary_content=response.text,
            search_references=[
                SearchReference(
                    site='lego',
                    url='lego',
                    content='lego',
                    title='lego',
                )  
            ]
        )


if __name__ == '__main__':
    # 测试查询
    test_queries = ["火箭积木列出", "怎样拼装屋顶"]
    
    # 实例化搜索引擎
    search_engine = BailianBotSearchEngine(
        bot_id='dbce22d58d7b4357b362b',
        api_key="sk-cacafff2c8d3446fbe6702adf4f"
    )
    
    # 执行搜索
    results = search_engine.search(test_queries)
    
    # 打印结果
    for result in results:
        print(f"查询: {result.query}")
        print(f"结果: {result.summary_content}")
        print("---")
